"""
LLM模型核心组件实现
包括注意力机制、前馈网络、位置编码等
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List
from model_config import LLMConfig


class RMSNorm(nn.Module):
    """RMS归一化层"""
    
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 计算RMS
        rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps)
        # 归一化并缩放
        return self.weight * (x / rms).type_as(x)


class SwiGLU(nn.Module):
    """SwiGLU激活函数"""
    
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播: SwiGLU(x) = Swish(W1(x)) * W2(x)"""
        gate = F.silu(self.gate_proj(x))  # Swish激活函数
        up = self.up_proj(x)
        return self.dropout(self.down_proj(gate * up))


def precompute_freqs_cis(dim: int, end: int, theta: float = 1e6) -> Tuple[torch.Tensor, torch.Tensor]:
    """预计算RoPE的频率矩阵"""
    # 计算频率
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 计算位置
    t = torch.arange(end, device=freqs.device)
    # 外积得到频率矩阵
    freqs = torch.outer(t, freqs).float()
    # 计算cos和sin
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin


def apply_rotary_pos_emb(
    q: torch.Tensor, 
    k: torch.Tensor, 
    cos: torch.Tensor, 
    sin: torch.Tensor,
    position_ids: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """应用旋转位置编码"""
    
    def rotate_half(x: torch.Tensor) -> torch.Tensor:
        """旋转一半维度"""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)
    
    # 确保cos和sin的维度与q, k匹配
    # cos和sin的形状应该是 [seq_len, head_dim]
    # q和k的形状是 [batch_size, seq_len, num_heads, head_dim]
    
    # 扩展cos和sin的维度以匹配q和k
    cos = cos.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, head_dim]
    sin = sin.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, head_dim]
    
    # 应用旋转位置编码
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    
    return q_embed, k_embed


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """重复KV头以匹配Q头数量"""
    if n_rep == 1:
        return x
    bs, slen, n_kv_heads, head_dim = x.shape
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )


class MultiHeadAttention(nn.Module):
    """多头注意力机制（支持GQA）"""
    
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.config = config
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_rep = config.num_rep
        
        # 投影层
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        
        # Dropout
        self.attn_dropout = nn.Dropout(config.attention_dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        
        # Flash Attention支持
        self.flash_attn = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        """前向传播"""
        bsz, seq_len, _ = hidden_states.shape
        
        # 计算Q, K, V
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)
        
        # 重塑为多头格式
        q = q.view(bsz, seq_len, self.num_heads, self.head_dim)
        k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)
        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim)
        
        # 应用旋转位置编码
        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos[:seq_len], sin[:seq_len])
        
        # KV缓存
        if past_key_value is not None:
            k = torch.cat([past_key_value[0], k], dim=1)
            v = torch.cat([past_key_value[1], v], dim=1)
        past_kv = (k, v) if use_cache else None
        
        # 转置为注意力计算格式
        q = q.transpose(1, 2)  # (bsz, num_heads, seq_len, head_dim)
        k = repeat_kv(k, self.num_rep).transpose(1, 2)
        v = repeat_kv(v, self.num_rep).transpose(1, 2)
        
        # 计算注意力
        if self.flash_attn and seq_len != 1:
            # 使用Flash Attention
            dropout_p = self.attn_dropout.p if self.training else 0.0
            attn_mask = None
            if attention_mask is not None:
                attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.num_heads, seq_len, -1)
                attn_mask = attn_mask.bool()
            
            attn_output = F.scaled_dot_product_attention(
                q, k, v, 
                attn_mask=attn_mask, 
                dropout_p=dropout_p, 
                is_causal=True
            )
        else:
            # 手动实现注意力
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            
            # 因果掩码
            causal_mask = torch.triu(
                torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
                diagonal=1
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)
            
            # 注意力掩码
            if attention_mask is not None:
                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
                extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
                scores = scores + extended_attention_mask
            
            # Softmax
            attn_weights = F.softmax(scores.float(), dim=-1).type_as(q)
            attn_weights = self.attn_dropout(attn_weights)
            attn_output = torch.matmul(attn_weights, v)
        
        # 输出投影
        attn_output = attn_output.transpose(1, 2).reshape(bsz, seq_len, -1)
        attn_output = self.resid_dropout(self.o_proj(attn_output))
        
        return attn_output, past_kv


class FeedForward(nn.Module):
    """前馈网络"""
    
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.dropout)
        self.act_fn = F.silu  # Swish激活函数
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播: SwiGLU(x) = Swish(W1(x)) * W2(x)"""
        return self.dropout(
            self.down_proj(
                self.act_fn(self.gate_proj(x)) * self.up_proj(x)
            )
        )


class MoEGate(nn.Module):
    """MoE门控网络"""
    
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.alpha = config.aux_loss_alpha
        self.seq_aux = config.seq_aux
        self.norm_topk_prob = config.norm_topk_prob
        
        # 门控权重
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))
        self.reset_parameters()
    
    def reset_parameters(self):
        """初始化参数"""
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
    
    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """前向传播"""
        bsz, seq_len, h = hidden_states.shape
        hidden_states = hidden_states.view(-1, h)
        
        # 计算门控分数
        logits = F.linear(hidden_states, self.weight, None)
        scores = F.softmax(logits, dim=-1)
        
        # Top-K选择
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        
        # 归一化权重
        if self.top_k > 1 and self.norm_topk_prob:
            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
            topk_weight = topk_weight / denominator
        
        # 计算辅助损失
        aux_loss = 0
        if self.training and self.alpha > 0.0:
            if self.seq_aux:
                # 序列级辅助损失
                scores_for_aux = scores.view(bsz, seq_len, -1)
                topk_idx_for_aux = topk_idx.view(bsz, -1)
                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
                ce.scatter_add_(
                    1, 
                    topk_idx_for_aux,
                    torch.ones(bsz, seq_len * self.top_k, device=hidden_states.device)
                ).div_(seq_len * self.top_k / self.n_routed_experts)
                aux_loss = (ce * scores_for_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
            else:
                # Token级辅助损失
                mask_ce = F.one_hot(topk_idx.view(-1), num_classes=self.n_routed_experts)
                ce = mask_ce.float().mean(0)
                Pi = scores.mean(0)
                fi = ce * self.n_routed_experts
                aux_loss = (Pi * fi).sum() * self.alpha
        
        return topk_idx, topk_weight, aux_loss


class MOEFeedForward(nn.Module):
    """MoE前馈网络"""
    
    def __init__(self, config: LLMConfig):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            FeedForward(config) for _ in range(config.n_routed_experts)
        ])
        self.gate = MoEGate(config)
        
        if config.n_shared_experts > 0:
            self.shared_experts = nn.ModuleList([
                FeedForward(config) for _ in range(config.n_shared_experts)
            ])
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        identity = x
        orig_shape = x.shape
        bsz, seq_len, _ = x.shape
        
        # 门控选择专家
        topk_idx, topk_weight, aux_loss = self.gate(x)
        x = x.view(-1, x.shape[-1])
        flat_topk_idx = topk_idx.view(-1)
        
        if self.training:
            # 训练时：重复输入并选择专家
            x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
            y = torch.empty_like(x, dtype=torch.float16)
            
            for i, expert in enumerate(self.experts):
                mask = flat_topk_idx == i
                if mask.any():
                    y[mask] = expert(x[mask]).to(y.dtype)
            
            # 加权求和
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            y = y.view(*orig_shape)
        else:
            # 推理时：优化计算
            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
        
        # 共享专家
        if self.config.n_shared_experts > 0:
            for expert in self.shared_experts:
                y = y + expert(identity)
        
        # 保存辅助损失
        self.aux_loss = aux_loss
        return y
    
    @torch.no_grad()
    def moe_infer(self, x: torch.Tensor, flat_expert_indices: torch.Tensor, flat_expert_weights: torch.Tensor) -> torch.Tensor:
        """MoE推理优化"""
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
        token_idxs = idxs // self.config.num_experts_per_tok
        
        for i, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
            if start_idx == end_idx:
                continue
            
            expert = self.experts[i]
            exp_token_idx = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idx]
            expert_out = expert(expert_tokens).to(expert_cache.dtype)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
        
        return expert_cache